{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 相似性检索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import Chroma\n",
    "from langchain.embeddings.dashscope import DashScopeEmbeddings\n",
    "from dotenv import load_dotenv\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "api_key = os.environ.get(\"DASHSCOPE_API_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_96004/692543272.py:5: LangChainDeprecationWarning: The class `Chroma` was deprecated in LangChain 0.2.9 and will be removed in 1.0. An updated version of the class exists in the :class:`~langchain-chroma package and should be used instead. To use it run `pip install -U :class:`~langchain-chroma` and import as `from :class:`~langchain_chroma import Chroma``.\n",
      "  vectordb_chinese = Chroma(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "29\n"
     ]
    }
   ],
   "source": [
    "persist_directory_chinese = \"./docs/chroma/matplotlib/\"\n",
    "\n",
    "embedding = DashScopeEmbeddings(dashscope_api_key=api_key)\n",
    "\n",
    "vectordb_chinese = Chroma(\n",
    "    persist_directory=persist_directory_chinese,\n",
    "    embedding_function=embedding\n",
    ")\n",
    "\n",
    "print(vectordb_chinese._collection.count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "texts_chinese = [\n",
    "    \"毒鹅膏菌（Amanita phalloides）具有大型且引人注目的地上（epigeous）子实体（basidiocarp）\",\n",
    "    \"一种具有大型子实体的蘑菇是毒鹅膏菌（Amanita phalloides）。某些品种全白。\",\n",
    "    \"A. phalloides, 又名死亡帽，是已知所有蘑菇中最有毒的一种。\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "smalldb_chinese = Chroma.from_texts(texts_chinese, embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "question_chinese = \"告诉我关于具有大型子实体的全白蘑菇的信息\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(metadata={}, page_content='一种具有大型子实体的蘑菇是毒鹅膏菌（Amanita phalloides）。某些品种全白。'),\n",
       " Document(metadata={}, page_content='毒鹅膏菌（Amanita phalloides）具有大型且引人注目的地上（epigeous）子实体（basidiocarp）')]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "smalldb_chinese.similarity_search(question_chinese, k=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 最大边际相关性：解决多样性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(metadata={}, page_content='一种具有大型子实体的蘑菇是毒鹅膏菌（Amanita phalloides）。某些品种全白。'),\n",
       " Document(metadata={}, page_content='A. phalloides, 又名死亡帽，是已知所有蘑菇中最有毒的一种。')]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "smalldb_chinese.max_marginal_relevance_search(question_chinese, k=2, fetch_k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "docs[0]:\n",
      "第⼀回： Matplotlib 初相识\n",
      "⼀、认识 matplotlib \n",
      "Matplotlib 是⼀个 Python 2D 绘图库，能够以多种硬拷⻉格式和跨平台的交互式环境⽣成出版物质量的图形，⽤来绘\n",
      "docs[1]:\n",
      "第⼀回： Matplotlib 初相识\n",
      "⼀、认识 matplotlib \n",
      "Matplotlib 是⼀个 Python 2D 绘图库，能够以多种硬拷⻉格式和跨平台的交互式环境⽣成出版物质量的图形，⽤来绘\n"
     ]
    }
   ],
   "source": [
    "question_chinese = \"Matplotlib是什么？\"\n",
    "docs_ss_chinese = vectordb_chinese.similarity_search(question_chinese, k=3)\n",
    "print(\"docs[0]:\")\n",
    "print(docs_ss_chinese[0].page_content[:100])\n",
    "print(\"docs[1]:\")\n",
    "print(docs_ss_chinese[1].page_content[:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "docs_mmr_chinese = vectordb_chinese.max_marginal_relevance_search(question_chinese,k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "第⼀回： Matplotlib 初相识\n",
      "⼀、认识 matplotlib \n",
      "Matplotlib 是⼀个 Python 2D 绘图库，能够以多种硬拷⻉格式和跨平台的交互式环境⽣成出版物质量的图形，⽤来绘\n"
     ]
    }
   ],
   "source": [
    "print(docs_mmr_chinese[0].page_content[:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2. patches\n",
      "matplotlib.patches.Patch 类是⼆维图形类，并且它是众多⼆维图形的⽗类，它的所有⼦类⻅matplotlib.patches API ，\n",
      "Patch 类的构造\n"
     ]
    }
   ],
   "source": [
    "print(docs_mmr_chinese[1].page_content[:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "By Datawhale 数据可视化开源⼩组\n",
      "© Copyright © Copyright 2021.\n",
      "⽤  np.random.randn(2, 150) ⽣成⼀组⼆维数据，使⽤两种⾮均匀⼦图的分\n"
     ]
    }
   ],
   "source": [
    "print(docs_mmr_chinese[2].page_content[:100])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 使用元数据解决特殊性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "question_chinese = \"他们在第二讲中对Figure说了些什么？\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "docs_chinese = vectordb_chinese.similarity_search(question_chinese,k=3,filter={\"source\":\"../data/第二回：艺术画笔见乾坤.pdf\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'creator': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36 Edg/141.0.0.0', 'page': 10, 'total_pages': 14, 'page_label': '11', 'producer': 'Skia/PDF m141', 'creationdate': '2025-10-07T13:45:33+00:00', 'source': '../data/第二回：艺术画笔见乾坤.pdf', 'title': '第二回：艺术画笔见乾坤 — fantastic-matplotlib', 'moddate': '2025-10-07T13:45:33+00:00'}\n",
      "{'creationdate': '2025-10-07T13:45:33+00:00', 'source': '../data/第二回：艺术画笔见乾坤.pdf', 'creator': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36 Edg/141.0.0.0', 'total_pages': 14, 'title': '第二回：艺术画笔见乾坤 — fantastic-matplotlib', 'page_label': '10', 'producer': 'Skia/PDF m141', 'moddate': '2025-10-07T13:45:33+00:00', 'page': 9}\n",
      "{'page': 12, 'title': '第二回：艺术画笔见乾坤 — fantastic-matplotlib', 'moddate': '2025-10-07T13:45:33+00:00', 'source': '../data/第二回：艺术画笔见乾坤.pdf', 'total_pages': 14, 'creator': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36 Edg/141.0.0.0', 'creationdate': '2025-10-07T13:45:33+00:00', 'page_label': '13', 'producer': 'Skia/PDF m141'}\n"
     ]
    }
   ],
   "source": [
    "for d in docs_chinese:\n",
    "    print(d.metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 在元数据中使用子查询检索器（LLM辅助检索）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms.tongyi import Tongyi\n",
    "from langchain.retrievers.self_query.base import SelfQueryRetriever\n",
    "from langchain.chains.query_constructor.base import AttributeInfo\n",
    "\n",
    "llm = Tongyi(api_key=api_key, model=\"qwen-flash\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata_field_info_chinese = [\n",
    "    AttributeInfo(\n",
    "        name=\"source\",\n",
    "        description=\"The lecture the chunk is from, should be one of `../data/第一回：Matplotlib初相识.pdf`,`../data/第二回：艺术画笔见乾坤.pdf`,or `../data/第三回：布局格式定方圆.pdf`\",\n",
    "        type=\"string\"\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"page\",\n",
    "        description=\"The page from the lecture\",\n",
    "        type=\"integer\"\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "document_content_description_chinese = \"Matplotlib 课堂讲义\"\n",
    "retriever_chinese = SelfQueryRetriever.from_llm(\n",
    "    llm,\n",
    "    vectordb_chinese,\n",
    "    document_content_description_chinese,\n",
    "    metadata_field_info_chinese,\n",
    "    verbose=True\n",
    ")\n",
    "question_chinese = \"他们在第二讲中对Figure做了些什么？\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "docs_chinese = retriever_chinese.get_relevant_documents(question_chinese)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "for d in docs_chinese:\n",
    "    print(d.metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(docs_chinese)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 压缩"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain.retrievers.document_compressors import LLMChainExtractor\n",
    "\n",
    "def pretty_print_docs(docs):\n",
    "    print(f\"\\n{'-'*100}\\n\".join([f\"Document {i+1}:\\n\\n\" + d.page_content for i, d in enumerate(docs)]))\n",
    "# 压缩器\n",
    "compressor = LLMChainExtractor.from_llm(llm)\n",
    "\n",
    "compression_retriver_chinese = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor,\n",
    "    base_retriever=vectordb_chinese.as_retriever(search_type='mmr',search_kwargs={'k':5,'fetch_k':50})\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document 1:\n",
      "\n",
      "Matplotlib 是一个 Python 2D 绘图库，能够以多种硬拷贝格式和跨平台的交互式环境生成出版物质量的图形，用来绘制各种静态，动态，交互式的图表。\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 2:\n",
      "\n",
      "Matplotlib 是一个用于创建高质量图表和可视化数据的 Python 绘图库。它提供了丰富的功能来绘制各种类型的图形，如折线图、散点图、柱状图、直方图、图像等。其核心组件包括：\n",
      "\n",
      "- **Axes**：坐标轴对象，用于绘图的主要区域。\n",
      "- **Artist Container**：容器类，用于管理图形元素（如线条、矩形、多边形等）。\n",
      "- **primitives**：基本图形元素，包括：\n",
      "  - **Line2D**：用于绘制二维曲线，支持实线、虚线、标记点等样式。\n",
      "  - **Rectangle**：用于绘制矩形，常见于柱状图、直方图等。\n",
      "  - **Polygon**：用于绘制多边形，常用于填充区域。\n",
      "  - **AxesImage**：用于显示图像数据。\n",
      "\n",
      "Matplotlib 的主要特点包括：\n",
      "- 支持多种绘图类型，如 `plot`（折线图）、`bar`（柱状图）、`scatter`（散点图）、`hist`（直方图）、`imshow`（图像显示）等。\n",
      "- 可通过参数灵活控制图形外观，如颜色、线宽、标记样式、线型等。\n",
      "- 提供了多种方式设置图形属性，包括在绘图函数中直接传参、通过对象方法修改或使用 `setp()` 函数。\n",
      "\n",
      "总之，Matplotlib 是 Python 中最广泛使用的数据可视化工具之一，适用于科学计算、数据分析和展示。\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 3:\n",
      "\n",
      "Matplotlib是Python中一个用于数据可视化的开源库。\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 4:\n",
      "\n",
      "matplotlib.patches.Wedge 类是楔型类。其基类是 matplotlib.patches.Patch ，它的构造函数：\n",
      "class matplotlib.patches.Wedge(center, r, theta1, theta2, width=None, **kwargs)\n",
      "⼀个 Wedge- 楔型  是以坐标 x,y 为中⼼，半径为 r ，从 θ1 扫到 θ2( 单位是度 ) 。\n",
      "如果宽度给定，则从内半径 r - 宽度到外半径 r 画出部分楔形。 wedge 中⽐较常⻅的是绘制饼状图。\n",
      "matplotlib.pyplot.pie 语法：\n",
      "matplotlib.pyplot.pie(x, explode=None, labels=None, colors=None, autopct=None, pctdistance=0.6,\n",
      "shadow=False, labeldistance=1.1, startangle=0, radius=1, counterclock=True, wedgeprops=None,\n",
      "textprops=None, center=0, 0, frame=False, rotatelabels=False, *, normalize=None, data=None)\n",
      "制作数据 x 的饼图，每个楔⼦的⾯积⽤ x/sum(x) 表⽰。\n"
     ]
    }
   ],
   "source": [
    "# 对源文档进行压缩\n",
    "question_chinese = \"Matplotlib是什么？\"\n",
    "compression_docs_chinese = compression_retriver_chinese.get_relevant_documents(question_chinese)\n",
    "pretty_print_docs(compression_docs_chinese)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
